{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609bc3f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _application.tester"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72449ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "import asyncio\n",
    "import collections\n",
    "import inspect\n",
    "from unittest.mock import AsyncMock, MagicMock\n",
    "import json\n",
    "from contextlib import asynccontextmanager\n",
    "from itertools import groupby\n",
    "from typing import *\n",
    "from types import ModuleType\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from fastkafka import KafkaEvent\n",
    "from fastkafka._application.app import FastKafka, AwaitedMock, _get_kafka_brokers\n",
    "from fastkafka._components.asyncapi import KafkaBroker, KafkaBrokers\n",
    "from fastkafka._components.helpers import unwrap_list_type\n",
    "from fastkafka._components.meta import delegates, export, patch\n",
    "from fastkafka._components.producer_decorator import unwrap_from_kafka_event\n",
    "from fastkafka._components.aiokafka_consumer_loop import ConsumeCallable\n",
    "from fastkafka._testing.apache_kafka_broker import ApacheKafkaBroker\n",
    "from fastkafka._testing.in_memory_broker import InMemoryBroker\n",
    "from fastkafka._testing.local_redpanda_broker import LocalRedpandaBroker\n",
    "from fastkafka._components.helpers import remove_suffix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ee08fec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "from pydantic import Field\n",
    "\n",
    "from fastkafka import EventMetadata, KafkaEvent\n",
    "from fastkafka._components.logger import get_logger, suppress_timestamps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d75f36a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: ok\n"
     ]
    }
   ],
   "source": [
    "suppress_timestamps()\n",
    "logger = get_logger(__name__, level=20)\n",
    "logger.info(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14650b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "# allows async calls in notebooks\n",
    "\n",
    "import nest_asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83484244",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca915b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "\n",
    "app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=\"9092\")))\n",
    "\n",
    "\n",
    "@app.consumes()\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "724f7abe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _get_broker_spec(bootstrap_server: str) -> KafkaBroker:\n",
    "    \"\"\"\n",
    "    Helper function to get the broker specification from the bootstrap server URL.\n",
    "\n",
    "    Args:\n",
    "        bootstrap_server: The bootstrap server URL in the format \"<host>:<port>\".\n",
    "\n",
    "    Returns:\n",
    "        A KafkaBroker object representing the broker specification.\n",
    "    \"\"\"\n",
    "    url = bootstrap_server.split(\":\")[0]\n",
    "    port = bootstrap_server.split(\":\")[1]\n",
    "    return KafkaBroker(url=url, port=port, description=\"\", protocol=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476ff67b",
   "metadata": {},
   "source": [
    "## Fastkafka Tester class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc14b119",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@export(\"fastkafka.testing\")\n",
    "class Tester(FastKafka):\n",
    "    __test__ = False\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        app: Union[FastKafka, List[FastKafka]],\n",
    "        *,\n",
    "        use_in_memory_broker: bool = True,\n",
    "    ):\n",
    "        \"\"\"Mirror-like object for testing a FastKafka application\n",
    "\n",
    "        Can be used as context manager\n",
    "\n",
    "        Args:\n",
    "            app: The FastKafka application to be tested.\n",
    "            use_in_memory_broker: Whether to use an in-memory broker for testing or not.\n",
    "        \"\"\"\n",
    "        self.apps = app if isinstance(app, list) else [app]\n",
    "\n",
    "        for app in self.apps:\n",
    "            app.create_mocks()\n",
    "\n",
    "        super().__init__()\n",
    "        self.mirrors: Dict[Any, Any] = {}\n",
    "        self._kafka_brokers = self.apps[0]._kafka_brokers\n",
    "        self._kafka_config[\"bootstrap_servers_id\"] = self.apps[0]._kafka_config[\n",
    "            \"bootstrap_servers_id\"\n",
    "        ]\n",
    "        self._create_mirrors()\n",
    "        self.use_in_memory_broker = use_in_memory_broker\n",
    "\n",
    "    async def _start_tester(self) -> None:\n",
    "        \"\"\"Starts the Tester\"\"\"\n",
    "        for app in self.apps:\n",
    "            await app.__aenter__()\n",
    "        self.create_mocks()\n",
    "        self._arrange_mirrors()\n",
    "        await super().__aenter__()\n",
    "        await asyncio.sleep(3)\n",
    "\n",
    "    async def _stop_tester(self) -> None:\n",
    "        \"\"\"Shuts down the Tester\"\"\"\n",
    "        await super().__aexit__(None, None, None)\n",
    "        for app in self.apps[::-1]:\n",
    "            await app.__aexit__(None, None, None)\n",
    "\n",
    "    def _create_mirrors(self) -> None:\n",
    "        pass\n",
    "\n",
    "    def _arrange_mirrors(self) -> None:\n",
    "        pass\n",
    "    \n",
    "    def _set_arguments_and_return_old(self,\n",
    "                      bootstrap_servers_id: Optional[str],\n",
    "                      use_in_memory_broker: bool\n",
    "                      ) -> Dict[Any, Any]:\n",
    "        initial_arguments: Dict[Any, Any] = dict()\n",
    "        initial_arguments[\"use_in_memory_broker\"] = self.use_in_memory_broker\n",
    "        self.use_in_memory_broker = use_in_memory_broker\n",
    "        \n",
    "        initial_arguments[\"bootstrap_servers_id\"] = self._kafka_config[\"bootstrap_servers_id\"]\n",
    "        if bootstrap_servers_id is None:\n",
    "            bootstrap_servers_id = self._kafka_config[\"bootstrap_servers_id\"]\n",
    "        else:\n",
    "            self._kafka_config[\"bootstrap_servers_id\"] = bootstrap_servers_id\n",
    "        \n",
    "        for app in self.apps:\n",
    "            initial_arguments[app] = app._kafka_config[\"bootstrap_servers_id\"]\n",
    "            app._kafka_config[\"bootstrap_servers_id\"] = bootstrap_servers_id\n",
    "            \n",
    "        return initial_arguments\n",
    "    \n",
    "    def _restore_initial_arguments(self,\n",
    "                                  initial_arguments: Dict[Any, Any]\n",
    "                                  ) -> None:\n",
    "        self.use_in_memory_broker = initial_arguments[\"use_in_memory_broker\"]\n",
    "        self._kafka_config[\"bootstrap_servers_id\"] = initial_arguments[\"bootstrap_servers_id\"]\n",
    "        \n",
    "        for app in self.apps:\n",
    "            app._kafka_config[\"bootstrap_servers_id\"] = initial_arguments[app]\n",
    "\n",
    "    @asynccontextmanager\n",
    "    async def using_external_broker(self,\n",
    "                                    bootstrap_servers_id: Optional[str] = None,\n",
    "                     ) -> AsyncGenerator[\"Tester\", None]:\n",
    "        \"\"\"Tester context manager for using external broker \n",
    "\n",
    "        Args:\n",
    "            bootstrap_servers_id: The bootstrap server of aplications.\n",
    "            \n",
    "        Returns:\n",
    "            self or None\n",
    "        \"\"\"\n",
    "        initial_arguments = self._set_arguments_and_return_old(bootstrap_servers_id, use_in_memory_broker=False)\n",
    "            \n",
    "        async with self._create_ctx() as ctx:\n",
    "            try:\n",
    "                yield self\n",
    "            finally:\n",
    "                self._restore_initial_arguments(initial_arguments)\n",
    "                \n",
    "                \n",
    "    @asynccontextmanager\n",
    "    async def using_inmemory_broker(self, \n",
    "                      bootstrap_servers_id: Optional[str] = None,\n",
    "                     ) -> AsyncGenerator[\"Tester\", None]:\n",
    "        \"\"\"Tester context manager for using in-memory broker \n",
    "\n",
    "        Args:\n",
    "            bootstrap_servers_id: The bootstrap server of aplications.\n",
    "            \n",
    "        Returns:\n",
    "            self or None\n",
    "        \"\"\"\n",
    "        initial_arguments = self._set_arguments_and_return_old(bootstrap_servers_id, use_in_memory_broker=True)\n",
    "            \n",
    "        async with self._create_ctx() as ctx:\n",
    "            try:\n",
    "                yield self\n",
    "            finally:\n",
    "                self._restore_initial_arguments(initial_arguments)\n",
    "            \n",
    "            \n",
    "    @asynccontextmanager\n",
    "    async def _create_ctx(self) -> AsyncGenerator[\"Tester\", None]:\n",
    "        if self.use_in_memory_broker == True:\n",
    "            with InMemoryBroker(): # type: ignore\n",
    "                await self._start_tester()\n",
    "                try:\n",
    "                    yield self\n",
    "                finally:\n",
    "                    await self._stop_tester()\n",
    "        else:\n",
    "            await self._start_tester()\n",
    "            try:\n",
    "                yield self\n",
    "            finally:\n",
    "                await self._stop_tester()\n",
    "\n",
    "    async def __aenter__(self) -> \"Tester\":\n",
    "        self._ctx = self._create_ctx()\n",
    "        return await self._ctx.__aenter__()\n",
    "\n",
    "    async def __aexit__(self, *args: Any) -> None:\n",
    "        await self._ctx.__aexit__(*args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63af0257",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "<ExceptionInfo RuntimeError('ok') tblen=1>\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "<ExceptionInfo RuntimeError('ok') tblen=1>\n"
     ]
    }
   ],
   "source": [
    "for _ in range(2):\n",
    "    with pytest.raises(RuntimeError) as e:\n",
    "        async with Tester(app) as tester:\n",
    "            assert tester.is_started\n",
    "            raise RuntimeError(\"ok\")\n",
    "\n",
    "    print(e)\n",
    "    assert not tester.is_started"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4de0c72a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'latest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Producing msg msg='signal'\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "tester = Tester(app)\n",
    "\n",
    "\n",
    "@tester.produces()\n",
    "async def to_preprocessed_signals(msg: TestMsg) -> TestMsg:\n",
    "    print(f\"Producing msg {msg}\")\n",
    "    return msg\n",
    "\n",
    "\n",
    "tester.to_preprocessed_signals = to_preprocessed_signals\n",
    "\n",
    "\n",
    "@tester.consumes(auto_offset_reset=\"latest\")\n",
    "async def on_predictions(msg: TestMsg):\n",
    "    pass\n",
    "\n",
    "\n",
    "async with tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await asyncio.sleep(5)\n",
    "    tester.mocks.on_predictions.assert_called()\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "938e5f89",
   "metadata": {},
   "source": [
    "## Test multiple brokers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0ce0e2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_1:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Producing msg msg='signal'\n",
      "Defined broker:  msg=TestMsg(msg='signal')\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "kafka_brokers_1 = dict(localhost=[dict(url=\"server_1\", port=\"9092\")])\n",
    "kafka_brokers_2 = dict(localhost=dict(url=\"server_2\", port=\"9092\"))\n",
    "\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed_signals\")\n",
    "async def on_preprocessed_signals_1(msg: TestMsg):\n",
    "    print(f\"Default broker:  {msg=}\")\n",
    "    \n",
    "@app.consumes(topic=\"preprocessed_signals\", brokers=kafka_brokers_2)\n",
    "async def on_preprocessed_signals_2(msg: TestMsg):\n",
    "    print(f\"Defined broker:  {msg=}\")\n",
    "\n",
    "tester = Tester(app)\n",
    "\n",
    "\n",
    "@tester.produces(brokers=kafka_brokers_2)\n",
    "async def to_preprocessed_signals(msg: TestMsg) -> TestMsg:\n",
    "    print(f\"Producing msg {msg}\")\n",
    "    return msg\n",
    "\n",
    "\n",
    "tester.to_preprocessed_signals = to_preprocessed_signals\n",
    "\n",
    "async with tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await asyncio.sleep(5)\n",
    "    await app.awaited_mocks.on_preprocessed_signals_2.assert_called(\n",
    "        timeout=5\n",
    "    )\n",
    "    await app.awaited_mocks.on_preprocessed_signals_1.assert_not_called(\n",
    "        timeout=5\n",
    "    )\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b69daca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_1:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': ['server_2:9092']}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_2:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_1:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Producing msg msg='signal'\n",
      "msg=TestMsg(msg='signal')\n",
      "Sending prediction: msg='prediction'\n",
      "tester: msg=TestMsg(msg='prediction')\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "kafka_brokers_1 = dict(localhost=dict(url=\"server_1\", port=\"9092\"))\n",
    "kafka_brokers_2 = dict(localhost=[dict(url=\"server_2\", port=\"9092\")])\n",
    "\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed_signals\", brokers=kafka_brokers_2)\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    print(f\"{msg=}\")\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "tester = Tester(app)\n",
    "\n",
    "\n",
    "@tester.produces(topic=\"preprocessed_signals\", brokers=kafka_brokers_2)\n",
    "async def to_preprocessed_signals(msg: TestMsg) -> TestMsg:\n",
    "    print(f\"Producing msg {msg}\")\n",
    "    return msg\n",
    "\n",
    "\n",
    "@tester.consumes(auto_offset_reset=\"earliest\", brokers=kafka_brokers_1)\n",
    "async def on_predictions(msg: TestMsg):\n",
    "    print(f\"tester: {msg=}\")\n",
    "\n",
    "\n",
    "tester.to_preprocessed_signals = to_preprocessed_signals\n",
    "\n",
    "async with tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await asyncio.sleep(5)\n",
    "    await app.awaited_mocks.on_preprocessed_signals.assert_called(timeout=5)\n",
    "    await tester.awaited_mocks.on_predictions.assert_called(timeout=5)\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ea8c473",
   "metadata": {},
   "source": [
    "## Mirroring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c20a5a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def mirror_producer(\n",
    "    topic: str, producer_f: Callable[..., Any], brokers: str, app: FastKafka\n",
    ") -> Callable[..., Any]:\n",
    "    \"\"\"\n",
    "    Decorator to create a mirrored producer function.\n",
    "\n",
    "    Args:\n",
    "        topic: The topic to produce to.\n",
    "        producer_f: The original producer function.\n",
    "        brokers: The brokers configuration.\n",
    "        app: The FastKafka application.\n",
    "\n",
    "    Returns:\n",
    "        The mirrored producer function.\n",
    "    \"\"\"\n",
    "    msg_type = inspect.signature(producer_f).return_annotation\n",
    "\n",
    "    msg_type_unwrapped = unwrap_list_type(unwrap_from_kafka_event(msg_type))\n",
    "\n",
    "    async def skeleton_func(msg: BaseModel) -> None:\n",
    "        pass\n",
    "\n",
    "    mirror_func = skeleton_func\n",
    "    sig = inspect.signature(skeleton_func)\n",
    "\n",
    "    # adjust name, take into consideration the origin app and brokers\n",
    "    # configuration so that we can differentiate those two\n",
    "    mirror_func.__name__ = f\"mirror_{id(app)}_on_{remove_suffix(topic).replace('.', '_').replace('-', '_')}_{abs(hash(brokers))}\"\n",
    "\n",
    "    # adjust arg and return val\n",
    "    sig = sig.replace(\n",
    "        parameters=[\n",
    "            inspect.Parameter(\n",
    "                name=\"msg\",\n",
    "                annotation=msg_type_unwrapped,\n",
    "                kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    mirror_func.__signature__ = sig  # type: ignore\n",
    "\n",
    "    return mirror_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e006e6d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_topic1() -> TestMsg:\n",
    "    pass\n",
    "\n",
    "\n",
    "@app.produces(topic=\"topic2\")\n",
    "async def some_log(in_var: int) -> TestMsg:\n",
    "    pass\n",
    "\n",
    "\n",
    "@app.produces(topic=\"topic2\", brokers=dict(localhost=dict(url=\"localhost\", port=9093)))\n",
    "async def some_log_1(in_var: int) -> TestMsg:\n",
    "    pass\n",
    "\n",
    "\n",
    "@app.produces(topic=\"topic2\", brokers=dict(localhost=dict(url=\"localhost\", port=9093)))\n",
    "async def some_log_2(in_var: int) -> TestMsg:\n",
    "    pass\n",
    "\n",
    "\n",
    "for topic, (producer_f, _, brokers, _) in app._producers_store.items():\n",
    "    mirror = mirror_producer(\n",
    "        topic,\n",
    "        producer_f,\n",
    "        brokers.model_dump_json() if brokers is not None else app._kafka_brokers.model_dump_json(),\n",
    "        app,\n",
    "    )\n",
    "    assert \"_\".join(mirror.__name__.split(\"_\")[2:-1]) == \"on_\" + remove_suffix(topic)\n",
    "    assert (\n",
    "        inspect.signature(mirror).parameters[\"msg\"].annotation.__name__\n",
    "        == inspect.Parameter(\n",
    "            name=\"msg\",\n",
    "            annotation=TestMsg,\n",
    "            kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,\n",
    "        ).annotation.__name__\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c20acfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)))\n",
    "\n",
    "@app.produces()\n",
    "async def to_topic1() -> TestMsg:\n",
    "    pass\n",
    "\n",
    "\n",
    "@app.produces(topic=\"topic2\")\n",
    "async def some_log(in_var: int) -> KafkaEvent[List[TestMsg]]:\n",
    "    pass\n",
    "\n",
    "\n",
    "for topic, (producer_f, _, brokers, _) in app._producers_store.items():\n",
    "    mirror = mirror_producer(\n",
    "        topic,\n",
    "        producer_f,\n",
    "        brokers.model_dump_json() if brokers is not None else app._kafka_brokers.model_dump_json(),\n",
    "        app\n",
    "    )\n",
    "    assert \"_\".join(mirror.__name__.split(\"_\")[2:-1]) == \"on_\" + remove_suffix(topic)\n",
    "    assert inspect.signature(mirror).parameters[\"msg\"].annotation.__name__ == \"TestMsg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e96d23ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def mirror_consumer(\n",
    "    topic: str, consumer_f: Callable[..., Any], brokers: str, app: FastKafka\n",
    ") -> Callable[[BaseModel], Coroutine[Any, Any, BaseModel]]:\n",
    "    \"\"\"\n",
    "    Decorator to create a mirrored consumer function.\n",
    "\n",
    "    Args:\n",
    "        topic: The topic to consume from.\n",
    "        consumer_f: The original consumer function.\n",
    "        brokers: The brokers configuration.\n",
    "        app: The FastKafka application.\n",
    "\n",
    "    Returns:\n",
    "        The mirrored consumer function.\n",
    "    \"\"\"\n",
    "    msg_type = inspect.signature(consumer_f).parameters[\"msg\"]\n",
    "\n",
    "    msg_type_unwrapped = unwrap_list_type(msg_type)\n",
    "\n",
    "    async def skeleton_func(msg: BaseModel) -> BaseModel:\n",
    "        return msg\n",
    "\n",
    "    mirror_func = skeleton_func\n",
    "    sig = inspect.signature(skeleton_func)\n",
    "\n",
    "    # adjust name, take into consideration the origin app and brokers\n",
    "    # configuration so that we can differentiate those two\n",
    "    mirror_func.__name__ = f\"mirror_{id(app)}_to_{remove_suffix(topic).replace('.', '_').replace('-', '_')}_{abs(hash(brokers))}\"\n",
    "\n",
    "    # adjust arg and return val\n",
    "    sig = sig.replace(\n",
    "        parameters=[msg_type], return_annotation=msg_type_unwrapped.annotation\n",
    "    )\n",
    "\n",
    "    mirror_func.__signature__ = sig  # type: ignore\n",
    "    return mirror_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79f408f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for topic, (consumer_f, _, _, brokers, _) in app._consumers_store.items():\n",
    "    mirror = mirror_consumer(\n",
    "        topic,\n",
    "        consumer_f,\n",
    "        brokers.model_dump_json() if brokers is not None else app._kafka_brokers.model_dump_json(),\n",
    "        app\n",
    "    )\n",
    "    assert \"_\".join(mirror.__name__.split(\"_\")[3:-1]) == \"to_\" + remove_suffix(topic)\n",
    "    assert (\n",
    "        inspect.signature(mirror).return_annotation.__name__ == TestMsg.__name__\n",
    "    ), inspect.signature(mirror).return_annotation.__name__\n",
    "    assert (\n",
    "        inspect.signature(mirror).parameters[\"msg\"].annotation.__name__\n",
    "        == inspect.Parameter(\n",
    "            name=\"msg\",\n",
    "            annotation=TestMsg,\n",
    "            kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,\n",
    "        ).annotation.__name__\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1737f993",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def _create_mirrors(self: Tester) -> None:\n",
    "    \"\"\"\n",
    "    Creates mirror functions for producers and consumers.\n",
    "\n",
    "    Iterates over the FastKafka application and its producers and consumers. For each consumer, it creates a mirror\n",
    "    consumer function using the `mirror_consumer` decorator. For each producer, it creates a mirror producer function\n",
    "    using the `mirror_producer` decorator. The mirror functions are stored in the `self.mirrors` dictionary and also\n",
    "    set as attributes on the Tester instance.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    for app in self.apps:\n",
    "        for topic, (consumer_f, _, _, brokers, _) in app._consumers_store.items():\n",
    "            mirror_f = mirror_consumer(\n",
    "                topic,\n",
    "                consumer_f,\n",
    "                brokers.model_dump_json() if brokers is not None else app._kafka_brokers.model_dump_json(),\n",
    "                app,\n",
    "            )\n",
    "            mirror_f = self.produces(  # type: ignore\n",
    "                topic=remove_suffix(topic),\n",
    "                brokers=brokers,\n",
    "            )(mirror_f)\n",
    "            self.mirrors[consumer_f] = mirror_f\n",
    "            setattr(self, mirror_f.__name__, mirror_f)\n",
    "        for topic, (producer_f, _, brokers, _) in app._producers_store.items():\n",
    "            mirror_f = mirror_producer(\n",
    "                topic,\n",
    "                producer_f,\n",
    "                brokers.model_dump_json() if brokers is not None else app._kafka_brokers.model_dump_json(),\n",
    "                app,\n",
    "            )\n",
    "            mirror_f = self.consumes(\n",
    "                topic=remove_suffix(topic),\n",
    "                brokers=brokers,\n",
    "            )(\n",
    "                mirror_f  # type: ignore\n",
    "            )\n",
    "            self.mirrors[producer_f] = mirror_f\n",
    "            setattr(self, mirror_f.__name__, mirror_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af0ccdda",
   "metadata": {},
   "outputs": [],
   "source": [
    "kafka_brokers_1 = dict(localhost=[dict(url=\"server_1\", port=9092)])\n",
    "kafka_brokers_2 = dict(localhost=dict(url=\"server_2\", port=9092))\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed_signals\", brokers=kafka_brokers_2)\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    print(f\"{msg=}\")\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "tester = Tester(app)\n",
    "\n",
    "assert hasattr(\n",
    "    tester,\n",
    "    f\"mirror_{id(app)}_to_preprocessed_signals_{abs(hash(_get_kafka_brokers(kafka_brokers_2).model_dump_json()))}\",\n",
    ")\n",
    "\n",
    "assert hasattr(\n",
    "    tester,\n",
    "    f\"mirror_{id(app)}_on_predictions_{abs(hash(_get_kafka_brokers(app._kafka_brokers).model_dump_json()))}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86bd35dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': ['server_1:9092']}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_1:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "msg=TestMsg(msg='signal')\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "kafka_brokers_1 = dict(localhost=[dict(url=\"server_1\", port=9092)])\n",
    "kafka_brokers_2 = dict(localhost=dict(url=\"server_2\", port=9092))\n",
    "\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed_signals\", brokers=kafka_brokers_2)\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    print(f\"{msg=}\")\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "async with Tester(app) as tester:\n",
    "    await getattr(\n",
    "        tester,\n",
    "        f\"mirror_{id(app)}_to_preprocessed_signals_{abs(hash(_get_kafka_brokers(kafka_brokers_2).model_dump_json()))}\",\n",
    "    )(TestMsg(msg=\"signal\"))\n",
    "    await asyncio.sleep(5)\n",
    "    await app.awaited_mocks.on_preprocessed_signals.assert_called(timeout=5)\n",
    "    await getattr(\n",
    "        tester.awaited_mocks,\n",
    "        f\"mirror_{id(app)}_on_predictions_{abs(hash(_get_kafka_brokers(kafka_brokers_1).model_dump_json()))}\",\n",
    "    ).assert_called(timeout=5)\n",
    "\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7af2c251",
   "metadata": {},
   "source": [
    "## Mirrors dict and syntax sugar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d5a5984",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class AmbiguousWarning:\n",
    "    \"\"\"\n",
    "    Warning class used for ambiguous topics.\n",
    "\n",
    "    Args:\n",
    "        topic: The ambiguous topic.\n",
    "        functions: List of function names associated with the ambiguous topic.\n",
    "    \"\"\"\n",
    "    def __init__(self, topic: str, functions: List[str]):\n",
    "        self.topic = topic\n",
    "        self.functions = functions\n",
    "\n",
    "    def __getattribute__(self, attr: str) -> Any:\n",
    "        raise RuntimeError(\n",
    "            f\"Ambiguous topic: {super().__getattribute__('topic')}, for functions: {super().__getattribute__('functions')}\\nUse Tester.mirrors[app.function] to resolve ambiguity\"\n",
    "        )\n",
    "\n",
    "    def __call__(self, *args: Any, **kwargs: Any) -> Any:\n",
    "        raise RuntimeError(\n",
    "            f\"Ambiguous topic: {self.topic}, for functions: {self.functions}\\nUse Tester.mirrors[app.function] to resolve ambiguity\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f86a0f07",
   "metadata": {},
   "outputs": [],
   "source": [
    "with pytest.raises(Exception) as e:\n",
    "    AmbiguousWarning(topic=\"some_topic\", functions=[\"some_functions\"])(TestMsg(msg=\"signal\"))\n",
    "    \n",
    "assert e.value.args[0] == \"Ambiguous topic: some_topic, for functions: ['some_functions']\\nUse Tester.mirrors[app.function] to resolve ambiguity\"\n",
    "\n",
    "with pytest.raises(Exception) as e:\n",
    "    AmbiguousWarning(topic=\"some_topic\", functions=[\"some_brokers\"]).assert_called(timeout=5)\n",
    "    \n",
    "assert e.value.args[0] == \"Ambiguous topic: some_topic, for functions: ['some_brokers']\\nUse Tester.mirrors[app.function] to resolve ambiguity\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a841f55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def set_sugar(\n",
    "    *,\n",
    "    tester: Tester,\n",
    "    prefix: str,\n",
    "    topic_brokers: Dict[str, Tuple[List[str], List[str]]],\n",
    "    topic: str,\n",
    "    brokers: str,\n",
    "    origin_function_name: str,\n",
    "    function: Callable[..., Union[Any, Awaitable[Any]]],\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Sets the sugar function for a topic.\n",
    "\n",
    "    Args:\n",
    "        tester: The Tester instance.\n",
    "        prefix: The prefix to use for the sugar function (e.g., \"to_\" or \"on_\").\n",
    "        topic_brokers: Dictionary to store the brokers and functions associated with each topic.\n",
    "        topic: The topic name.\n",
    "        brokers: The brokers configuration.\n",
    "        origin_function_name: The name of the original function.\n",
    "        function: The mirror function to be set as the sugar function.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    brokers_for_topic, functions_for_topic = topic_brokers.get(topic, ([], []))\n",
    "    if brokers not in brokers_for_topic:\n",
    "        brokers_for_topic.append(brokers)\n",
    "        functions_for_topic.append(origin_function_name)\n",
    "        topic_brokers[topic] = (brokers_for_topic, functions_for_topic)\n",
    "    if len(brokers_for_topic) == 1:\n",
    "        setattr(tester, f\"{prefix}{topic}\", function)\n",
    "    else:\n",
    "        setattr(\n",
    "            tester, f\"{prefix}{topic}\", AmbiguousWarning(topic, functions_for_topic)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d8480ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def _arrange_mirrors(self: Tester) -> None:\n",
    "    \"\"\"\n",
    "    Arranges the mirror functions.\n",
    "\n",
    "    Iterates over the FastKafka application and its producers and consumers. For each consumer, it retrieves the mirror\n",
    "    function from the `self.mirrors` dictionary and sets it as an attribute on the Tester instance. It also sets the\n",
    "    sugar function using the `set_sugar` function. For each producer, it retrieves the mirror function and sets it as\n",
    "    an attribute on the Tester instance. It also sets the sugar function for the awaited mocks. Finally, it creates the\n",
    "    `mocks` and `awaited_mocks` namedtuples and sets them as attributes on the Tester instance.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    topic_brokers: Dict[str, Tuple[List[str], List[str]]] = {}\n",
    "    mocks = {}\n",
    "    awaited_mocks = {}\n",
    "    for app in self.apps:\n",
    "        for topic, (consumer_f, _, _, brokers, _) in app._consumers_store.items():\n",
    "            mirror_f = self.mirrors[consumer_f]\n",
    "            self.mirrors[getattr(app, consumer_f.__name__)] = mirror_f\n",
    "            set_sugar(\n",
    "                tester=self,\n",
    "                prefix=\"to_\",\n",
    "                topic_brokers=topic_brokers,\n",
    "                topic=remove_suffix(topic).replace(\".\", \"_\").replace(\"-\", \"_\"),\n",
    "                brokers=brokers.model_dump_json()\n",
    "                if brokers is not None\n",
    "                else app._kafka_brokers.model_dump_json(),\n",
    "                origin_function_name=consumer_f.__name__,\n",
    "                function=mirror_f,\n",
    "            )\n",
    "\n",
    "            mocks[\n",
    "                f\"to_{remove_suffix(topic).replace('.', '_').replace('-', '_')}\"\n",
    "            ] = getattr(self.mocks, mirror_f.__name__)\n",
    "            awaited_mocks[\n",
    "                f\"to_{remove_suffix(topic).replace('.', '_').replace('-', '_')}\"\n",
    "            ] = getattr(self.awaited_mocks, mirror_f.__name__)\n",
    "\n",
    "        for topic, (producer_f, _, brokers, _) in app._producers_store.items():\n",
    "            mirror_f = self.mirrors[producer_f]\n",
    "            self.mirrors[getattr(app, producer_f.__name__)] = getattr(\n",
    "                self.awaited_mocks, mirror_f.__name__\n",
    "            )\n",
    "            set_sugar(\n",
    "                tester=self,\n",
    "                prefix=\"on_\",\n",
    "                topic_brokers=topic_brokers,\n",
    "                topic=remove_suffix(topic).replace(\".\", \"_\").replace(\"-\", \"_\"),\n",
    "                brokers=brokers.model_dump_json()\n",
    "                if brokers is not None\n",
    "                else app._kafka_brokers.model_dump_json(),\n",
    "                origin_function_name=producer_f.__name__,\n",
    "                function=getattr(self.awaited_mocks, mirror_f.__name__),\n",
    "            )\n",
    "            mocks[\n",
    "                f\"on_{remove_suffix(topic).replace('.', '_').replace('-', '_')}\"\n",
    "            ] = getattr(self.mocks, mirror_f.__name__)\n",
    "            awaited_mocks[\n",
    "                f\"on_{remove_suffix(topic).replace('.', '_').replace('-', '_')}\"\n",
    "            ] = getattr(self.awaited_mocks, mirror_f.__name__)\n",
    "\n",
    "    AppMocks = collections.namedtuple(  # type: ignore\n",
    "        f\"{self.__class__.__name__}Mocks\", [f_name for f_name in mocks]\n",
    "    )\n",
    "    setattr(self, \"mocks\", AppMocks(**mocks))\n",
    "    setattr(self, \"awaited_mocks\", AppMocks(**awaited_mocks))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0cb8429",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['store_product']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "# Test mock calls immutability\n",
    "\n",
    "class Currency(BaseModel):\n",
    "    currency: str = Field(..., description=\"Currency\")\n",
    "\n",
    "app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"server_2\", port=9092)))\n",
    "\n",
    "@app.consumes(prefix=\"on\", topic=\"store_product\")\n",
    "async def on_store_product(msg: Currency):\n",
    "    msg.currency = \"EUR\"\n",
    "    \n",
    "async with Tester(app).using_inmemory_broker() as tester:\n",
    "    await tester.to_store_product(Currency(currency=\"HRK\"))\n",
    "    await app.awaited_mocks.on_store_product.assert_called_with(\n",
    "        Currency(currency=\"HRK\"),\n",
    "        timeout=5\n",
    "    )\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d84884",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "# Test batch mirroring\n",
    "\n",
    "\n",
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "\n",
    "second_app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)))\n",
    "\n",
    "\n",
    "@second_app.consumes()\n",
    "async def on_preprocessed_signals(msg: TestMsg, meta: EventMetadata):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@second_app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> List[TestMsg]:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return [prediction]\n",
    "\n",
    "\n",
    "async with Tester(second_app) as tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await tester.awaited_mocks.on_predictions.assert_called(timeout=5)\n",
    "    tester.mocks.on_predictions.assert_called()\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b59f70e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_1:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_1:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "msg=TestMsg(msg='signal')\n",
      "Sending prediction: msg='prediction'\n",
      "[WARNING] fastkafka._components.task_streaming: e=AttributeError(\"'NoneType' object has no attribute 'create_batch'\")\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "kafka_brokers_1 = dict(localhost=dict(url=\"server_1\", port=9092))\n",
    "kafka_brokers_2 = dict(localhost=dict(url=\"server_2\", port=9092))\n",
    "\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed_signals\")\n",
    "async def on_preprocessed_signals_1(msg: TestMsg):\n",
    "    print(f\"{msg=}\")\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed_signals\", brokers=kafka_brokers_2)\n",
    "async def on_preprocessed_signals_2(msg: TestMsg):\n",
    "    print(f\"{msg=}\")\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "async with Tester(app) as tester:\n",
    "    with pytest.raises(Exception) as exception_produce:\n",
    "        await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    assert (\n",
    "        exception_produce.value.args[0]\n",
    "        == \"Ambiguous topic: preprocessed_signals, for functions: ['on_preprocessed_signals_1', 'on_preprocessed_signals_2']\\nUse Tester.mirrors[app.function] to resolve ambiguity\"\n",
    "    )\n",
    "    await tester.mirrors[on_preprocessed_signals_1](TestMsg(msg=\"signal\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ed56429",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': ['server_2:9092']}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': ['server_2:9092']}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_2:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed-signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_2:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "# Test topics with hyphens\n",
    "\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.consumes(topic=\"preprocessed-signals\")\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces(topic=\"predictions\")\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "async with Tester(app) as tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "    await tester.awaited_mocks.on_predictions.assert_called_with(\n",
    "        TestMsg(msg=\"prediction\"),\n",
    "        timeout=60,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63606cfb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': ['server_2:9092']}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_2:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "kafka_brokers_1 = dict(localhost=[dict(url=\"server_1\", port=9092)])\n",
    "kafka_brokers_1 = dict(localhost=[dict(url=\"server_2\", port=9092)])\n",
    "\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "\n",
    "@app.produces(topic=\"predictions\")\n",
    "async def to_predictions_1(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "@app.produces(topic=\"predictions\", brokers=kafka_brokers_2)\n",
    "async def to_predictions_2(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "async with Tester(app) as tester:\n",
    "    with pytest.raises(Exception) as exception_consume:\n",
    "        await tester.on_predictions.assert_called(timeout=5)\n",
    "    assert (\n",
    "        exception_consume.value.args[0]\n",
    "        == \"Ambiguous topic: predictions, for functions: ['to_predictions_1', 'to_predictions_2']\\nUse Tester.mirrors[app.function] to resolve ambiguity\"\n",
    "    )\n",
    "    await tester.mirrors[app.to_predictions_1].assert_not_called(timeout=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c10ea363",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'localhost:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "# Test KafkaEvent mirroring and consumer batching\n",
    "\n",
    "\n",
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "\n",
    "second_app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"localhost\", port=9092)))\n",
    "\n",
    "\n",
    "@second_app.consumes()\n",
    "async def on_preprocessed_signals(msg: List[TestMsg]):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@second_app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> KafkaEvent[TestMsg]:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return KafkaEvent(message=prediction, key=b\"123\")\n",
    "\n",
    "\n",
    "async with Tester(second_app) as tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await tester.on_predictions.assert_called(timeout=5)\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f237621b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': ['server_2:9092']}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server_2:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': ['server_2:9092']}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server_2:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "# Initiate tester with two apps\n",
    "\n",
    "\n",
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "\n",
    "second_app = FastKafka(kafka_brokers=dict(localhost=dict(url=\"server_2\", port=9092)))\n",
    "\n",
    "\n",
    "@second_app.consumes()\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@second_app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "async with Tester([app, second_app]) as tester:\n",
    "    await tester.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await tester.mirrors[second_app.to_predictions].assert_called(timeout=5)\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aef439e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestMsg(BaseModel):\n",
    "    msg: str = Field(...)\n",
    "\n",
    "kafka_brokers_1 = dict(localhost=dict(url=\"server\", port=9092), production=dict(url=\"prod_server\", port=9092))\n",
    "kafka_brokers_2 = dict(localhost=dict(url=\"broker_2_server\", port=9092), production=dict(url=\"broker_2_prod_server\", port=9092))\n",
    "app = FastKafka(kafka_brokers=kafka_brokers_1)\n",
    "\n",
    "@app.consumes()\n",
    "async def on_preprocessed_signals(msg: TestMsg):\n",
    "    await to_predictions(TestMsg(msg=\"prediction\"))\n",
    "\n",
    "\n",
    "@app.produces()\n",
    "async def to_predictions(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction\n",
    "\n",
    "\n",
    "@app.produces(topic=\"predictions\", brokers=kafka_brokers_2)\n",
    "async def to_predictions_2(prediction: TestMsg) -> TestMsg:\n",
    "    print(f\"Sending prediction: {prediction}\")\n",
    "    return prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e524f7c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'broker_2_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'broker_2_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'prod_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'broker_2_prod_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'prod_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'prod_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'prod_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'broker_2_prod_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "tester = Tester([app])\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    async with tester.using_external_broker() as tester1:\n",
    "        assert tester1.use_in_memory_broker == False\n",
    "        bootstraps_topics_groups = list(broker.topic_groups.keys())\n",
    "        # kafka_brokers1 localhost\n",
    "        assert bootstraps_topics_groups[0][0] == \"server:9092\"\n",
    "        assert bootstraps_topics_groups[1][0] == \"server:9092\"\n",
    "        # kafka_brokers2 localhost\n",
    "        assert bootstraps_topics_groups[2][0] == \"broker_2_server:9092\"\n",
    "        assert app._kafka_config[\"bootstrap_servers_id\"] == 'localhost'\n",
    "        assert tester1._kafka_config[\"bootstrap_servers_id\"] == app._kafka_config[\"bootstrap_servers_id\"]\n",
    "\n",
    "assert tester.use_in_memory_broker == True\n",
    "        \n",
    "with InMemoryBroker() as broker:   \n",
    "    async with tester.using_external_broker(bootstrap_servers_id=\"production\") as tester2:\n",
    "        assert tester2.use_in_memory_broker == False\n",
    "        bootstraps_topics_groups = list(broker.topic_groups.keys())\n",
    "        # kafka_brokers1 production\n",
    "        assert bootstraps_topics_groups[0][0] == \"prod_server:9092\"\n",
    "        assert bootstraps_topics_groups[1][0] == \"prod_server:9092\"\n",
    "        # kafka_brokers2 production\n",
    "        assert bootstraps_topics_groups[2][0] == \"broker_2_prod_server:9092\"\n",
    "        \n",
    "        assert app._kafka_config[\"bootstrap_servers_id\"] == 'production'\n",
    "        assert tester2._kafka_config[\"bootstrap_servers_id\"] == app._kafka_config[\"bootstrap_servers_id\"]\n",
    "        \n",
    "assert tester.use_in_memory_broker == True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fba16412",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'broker_2_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'broker_2_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'prod_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'broker_2_prod_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'prod_server:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'prod_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['preprocessed_signals']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'prod_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'auto_offset_reset': 'earliest', 'max_poll_records': 100, 'bootstrap_servers': 'broker_2_prod_server:9092'}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "Sending prediction: msg='prediction'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "tester = Tester([app])\n",
    "\n",
    "async with tester.using_inmemory_broker() as tester1:\n",
    "    assert tester1.use_in_memory_broker == True\n",
    "    await tester1.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await tester1.mirrors[app.to_predictions].assert_called(timeout=5)\n",
    "\n",
    "    assert app._kafka_config[\"bootstrap_servers_id\"] == 'localhost'\n",
    "    assert tester1._kafka_config[\"bootstrap_servers_id\"] == app._kafka_config[\"bootstrap_servers_id\"]\n",
    "\n",
    "assert tester.use_in_memory_broker == True\n",
    "\n",
    "async with tester.using_inmemory_broker(bootstrap_servers_id=\"production\") as tester2:\n",
    "    assert tester2.use_in_memory_broker == True\n",
    "    await tester2.to_preprocessed_signals(TestMsg(msg=\"signal\"))\n",
    "    await tester2.mirrors[app.to_predictions].assert_called(timeout=5)\n",
    "\n",
    "    assert app._kafka_config[\"bootstrap_servers_id\"] == 'production'\n",
    "    assert tester2._kafka_config[\"bootstrap_servers_id\"] == app._kafka_config[\"bootstrap_servers_id\"]\n",
    "    \n",
    "assert tester.use_in_memory_broker == True"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
